import paddle

class FeedFroward(paddle.nn.Layer):
    def __init__(self, hidden_dim):
        super(FeedFroward, self).__init__()
        self.fc1 = paddle.nn.Linear(hidden_dim, hidden_dim // 2, bias_attr=False)
        self.fc2 = paddle.nn.Linear(hidden_dim // 2, hidden_dim, bias_attr=False)

    def forward(self, feed_x):
        feed_x = self.fc1(feed_x)
        # sin cos 一样 其次是silu 和tanh
        feed_x = paddle.sin(feed_x)
        feed_x = self.fc2(feed_x)
        return feed_x


class CVFroward(paddle.nn.Layer):
    def __init__(self, hidden_dim, k_size):
        super(CVFroward, self).__init__()
        self.fc1 = paddle.nn.Conv1D(in_channels=hidden_dim, out_channels=hidden_dim // 2, bias_attr=False,
                                    kernel_size=k_size, stride=1)
        self.fc2 = paddle.nn.Conv1D(in_channels=hidden_dim // 2, out_channels=hidden_dim, bias_attr=False,
                                    kernel_size=k_size, stride=1)

        self.pad = k_size - 1

    def forward(self, feed_x):
        feed_x = feed_x.transpose([0, 2, 1])
        zero_pad = paddle.zeros(shape=[feed_x.shape[0], feed_x.shape[1], self.pad], dtype=feed_x.dtype)
        feed_x = paddle.concat([zero_pad, feed_x], -1)
        feed_x = self.fc1(feed_x)
        # sin cos 一样 先是silu 和tanh
        feed_x = paddle.sin(feed_x)
        zero_pad = paddle.zeros(shape=[feed_x.shape[0], feed_x.shape[1], self.pad], dtype=feed_x.dtype)

        feed_x = paddle.concat([zero_pad, feed_x], -1)
        feed_x = self.fc2(feed_x)
        feed_x = feed_x.transpose([0, 2, 1])

        return feed_x


class ImEncoder(paddle.nn.Layer):
    def __init__(self):
        super(ImEncoder, self).__init__()
        self.cva = paddle.nn.Conv2D(3, 6, kernel_size=(3, 3), padding=0, stride=2, bias_attr=False)
        self.cvb = paddle.nn.Conv2D(6, 12, kernel_size=(3, 3), padding=0, stride=2, bias_attr=False)
        self.cvc = paddle.nn.Conv2D(12, 6, kernel_size=(3, 3), padding=0, stride=2, bias_attr=False)
        self.cvd = paddle.nn.Conv2D(6, 3, kernel_size=(3, 3), padding=0, stride=2, bias_attr=False)
        self.cve = paddle.nn.Conv2D(3, 3, kernel_size=(3, 3), padding=0, stride=2, bias_attr=False)
        self.layer_nor_a = paddle.nn.LayerNorm(127)
        self.layer_nor_b = paddle.nn.LayerNorm(63)
        self.layer_nor_c = paddle.nn.LayerNorm(31)
        self.layer_nor_d = paddle.nn.LayerNorm(15)
        self.layer_nor_e = paddle.nn.LayerNorm(7)
        self.cv_a = CVFroward(127, 3)
        self.cv_b = CVFroward(63, 3)
        self.cv_c = CVFroward(31, 3)
        self.cv_d = CVFroward(15, 3)
        self.cv_e = CVFroward(7, 3)
        self.ff_a = FeedFroward(127)
        self.ff_b = FeedFroward(63)
        self.ff_c = FeedFroward(31)
        self.ff_d = FeedFroward(15)
        self.ff_e = FeedFroward(7)

    def forward(self, x):
        x = self.cva(x)  # 1 6 127,127
        x = x.reshape([x.shape[0], -1, x.shape[-1]])
        x += self.cv_a(x)
        x = self.layer_nor_a(x)
        x += self.ff_a(x)
        x = self.layer_nor_a(x)
        x = x.reshape([x.shape[0], -1, x.shape[-1], x.shape[-1]])

        x = self.cvb(x)  # 1 12 63,63
        x = x.reshape([x.shape[0], -1, x.shape[-1]])
        x += self.cv_b(x)
        x = self.layer_nor_b(x)
        x += self.ff_b(x)
        x = self.layer_nor_b(x)
        x = x.reshape([x.shape[0], -1, x.shape[-1], x.shape[-1]])

        x = self.cvc(x)  # 1 6 31,31
        x = x.reshape([x.shape[0], -1, x.shape[-1]])
        x += self.cv_c(x)
        x = self.layer_nor_c(x)
        x += self.ff_c(x)
        x = self.layer_nor_c(x)
        x = x.reshape([x.shape[0], -1, x.shape[-1], x.shape[-1]])

        x = self.cvd(x)  # 1 3 15,15
        x = x.reshape([x.shape[0], -1, x.shape[-1]])
        x += self.cv_d(x)
        x = self.layer_nor_d(x)
        x += self.ff_d(x)
        x = self.layer_nor_d(x)
        x = x.reshape([x.shape[0], -1, x.shape[-1], x.shape[-1]])

        x = self.cve(x)  # 1 3 7,7
        x = x.reshape([x.shape[0], -1, x.shape[-1]])
        x += self.cv_e(x)
        x = self.layer_nor_e(x)
        x += self.ff_e(x)
        x = self.layer_nor_e(x)
        x = x.reshape([x.shape[0], -1, x.shape[-1], x.shape[-1]])

        return x


class ImDecoder(paddle.nn.Layer):
    def __init__(self):
        super(ImDecoder, self).__init__()
        self.cve = paddle.nn.Conv2DTranspose(6, 3, kernel_size=(3, 3), padding=0, stride=2, bias_attr=False)
        self.cvd = paddle.nn.Conv2DTranspose(12, 6, kernel_size=(3, 3), padding=0, stride=2, bias_attr=False)
        self.cvc = paddle.nn.Conv2DTranspose(6, 12, kernel_size=(3, 3), padding=0, stride=2, bias_attr=False)
        self.cvb = paddle.nn.Conv2DTranspose(3, 6, kernel_size=(3, 3), padding=0, stride=2, bias_attr=False)
        self.cva = paddle.nn.Conv2DTranspose(3, 3, kernel_size=(3, 3), padding=0, stride=2, bias_attr=False)
        self.layer_nor_e = paddle.nn.LayerNorm(127)
        self.layer_nor_d = paddle.nn.LayerNorm(63)
        self.layer_nor_c = paddle.nn.LayerNorm(31)
        self.layer_nor_b = paddle.nn.LayerNorm(15)

        self.cv_e = CVFroward(127, 3)
        self.cv_d = CVFroward(63, 3)
        self.cv_c = CVFroward(31, 3)
        self.cv_b = CVFroward(15, 3)
        self.cv_a = CVFroward(7, 3)
        self.ff_e = FeedFroward(127)
        self.ff_d = FeedFroward(63)
        self.ff_c = FeedFroward(31)
        self.ff_b = FeedFroward(15)

    def forward(self, x):
        x = self.cva(x)  # 1 6 127,127
        x = x.reshape([x.shape[0], -1, x.shape[-1]])
        x += self.cv_b(x)
        x = self.layer_nor_b(x)
        x += self.ff_b(x)
        x = self.layer_nor_b(x)
        x = x.reshape([x.shape[0], -1, x.shape[-1], x.shape[-1]])

        x = self.cvb(x)  # 1 12 63,63
        x = x.reshape([x.shape[0], -1, x.shape[-1]])
        x += self.cv_c(x)
        x = self.layer_nor_c(x)
        x += self.ff_c(x)
        x = self.layer_nor_c(x)
        x = x.reshape([x.shape[0], -1, x.shape[-1], x.shape[-1]])

        x = self.cvc(x)  # 1 6 31,31
        x = x.reshape([x.shape[0], -1, x.shape[-1]])
        x += self.cv_d(x)
        x = self.layer_nor_d(x)
        x += self.ff_d(x)
        x = self.layer_nor_d(x)
        x = x.reshape([x.shape[0], -1, x.shape[-1], x.shape[-1]])

        x = self.cvd(x)  # 1 3 15,15
        x = x.reshape([x.shape[0], -1, x.shape[-1]])
        x += self.cv_e(x)
        x = self.layer_nor_e(x)
        x += self.ff_e(x)
        x = self.layer_nor_e(x)
        x = x.reshape([x.shape[0], -1, x.shape[-1], x.shape[-1]])
        x = self.cve(x)

        return x


class ImEncoderDecoder(paddle.nn.Layer):
    def __init__(self):
        super(ImEncoderDecoder, self).__init__()
        self.en = ImEncoder()
        self.de = ImDecoder()

    def forward(self, feed_x):
        en = self.en(feed_x)
        feed_x = self.de(en)
        return feed_x, en


if __name__ == '__main__':
    net = ImEncoderDecoder()
    loss_f = paddle.nn.MSELoss()
    optimizer = paddle.optimizer.Adam(learning_rate=0.0001, parameters=net.parameters())

    for i in range(100):
        image = paddle.randn([10, 3, 256, 256])
        out, en = net(image)
        loss = loss_f(image[:, :, 1:, 1:], out)
        print(loss.item())
        optimizer.clear_grad()
        loss.backward()
        optimizer.step()
